
from langchain_community.vectorstores import Chroma
from langchain_openai.embeddings import AzureOpenAIEmbeddings  # 导入嵌入模型
from langchain_community.document_loaders import PyPDFLoader
from tool import get_azure_endpoint, get_api_key, get_api_version
if __name__ == '__main__':
    embedding = AzureOpenAIEmbeddings(
        azure_endpoint=get_azure_endpoint().rstrip('/'),  # 移除尾部斜杠，只保留基础URL
        model="text-embedding-3-small",  # 重命名为 azure_deployment
        api_key=get_api_key(),
        api_version=get_api_version()
    )
    vectordb_chinese = Chroma(
        persist_directory="./docs/chroma/matplotlib",
        embedding_function=embedding,
    )
    question_chinese = "Matplotlib是什么?"

    docs_chinese = vectordb_chinese.similarity_search(question_chinese, k=5)

    print("docs[0]")
    print(docs_chinese[0])
    print("docs[1]")
    print(docs_chinese[1])

    question_chinese = "他们在第二讲中对Figure说了些什么?"
    docs_chinese = vectordb_chinese.similarity_search(question_chinese, k=5)
    for doc_chinese in docs_chinese:
        print(doc_chinese.metadata)